package core

import (
	"container/heap"
	"sync"
)

type Status int

const (
	StatusSent Status = iota
	StatusLoss
	StatusTimeExceed
	StatusReachDst
	StatusEchoReply
	StatusRST
	StatusSYNACK
)

type PktRecord struct {
	Key string

	Turn    int
	Device  string
	SrcIP   string
	DstIP   string
	SrcPort uint16
	DstPort uint16
	TTL     uint8

	TxTs     int64
	RxTs     int64
	Received bool
	ReachDst bool
	Timeout  bool
	// Status   Status
}

type FlowRecord struct {
	Key string

	Turn int

	SrcIP      string
	DstIP      string
	SrcPort    uint16
	DstPort    uint16
	PktRecords []*PktRecord
	Text       string
}

type TraceHopStat struct {
	DeviceSet map[string]bool

	Total      int
	Loss       int
	MinLatency int64
	AvgLatency int64
	MaxLatency int64
}

type TraceStat struct {
	PktRecordBucket       PktRecordBucket
	AbnormalFlowRecordMap map[string][]*FlowRecord // 记录存在丢包(无响应)的trace信息,key为五元组
	TraceHopStats         []*TraceHopStat          // 每一跳的数据统计
}

func NewTraceStat(maxTTL uint8) *TraceStat {
	return &TraceStat{
		PktRecordBucket:       PktRecordBucket{PktRecoredMap: make(map[string]*PktRecord)},
		AbnormalFlowRecordMap: make(map[string][]*FlowRecord),
		TraceHopStats:         make([]*TraceHopStat, maxTTL+1),
	}
}

func DeepCopyFlowRecord(origin *FlowRecord) *FlowRecord {
	copied := &FlowRecord{
		Key:     origin.Key,
		Turn:    origin.Turn,
		SrcIP:   origin.SrcIP,
		SrcPort: origin.SrcPort,
		DstIP:   origin.DstIP,
		DstPort: origin.DstPort,
	}

	copied.PktRecords = append(copied.PktRecords, origin.PktRecords...)
	return copied
}

type PktRecordBucket struct {
	Mu            sync.RWMutex
	PktRecordHeap PktRecordHeap
	PktRecoredMap map[string]*PktRecord
}

func (b *PktRecordBucket) Add(pktRecord *PktRecord) {
	b.Mu.Lock()
	defer b.Mu.Unlock()
	b.PktRecoredMap[pktRecord.Key] = pktRecord
	heap.Push(&b.PktRecordHeap, pktRecord)
}

func (b *PktRecordBucket) Update(pktRecord *PktRecord) *PktRecord {
	b.Mu.Lock()
	defer b.Mu.Unlock()
	toUpdate := b.PktRecoredMap[pktRecord.Key]
	if toUpdate != nil {
		toUpdate.Received = pktRecord.Received
		toUpdate.ReachDst = pktRecord.ReachDst
		toUpdate.RxTs = pktRecord.RxTs
		toUpdate.Device = pktRecord.Device
	}

	return toUpdate
}

func (b *PktRecordBucket) Pop() *PktRecord {
	b.Mu.Lock()
	defer b.Mu.Unlock()
	v := heap.Pop(&b.PktRecordHeap).(*PktRecord)
	delete(b.PktRecoredMap, v.Key)
	return v
}

func (b *PktRecordBucket) Peek() *PktRecord {
	b.Mu.Lock()
	defer b.Mu.Unlock()

	if len(b.PktRecordHeap) == 0 {
		return nil
	}

	return b.PktRecordHeap[0]
}

func (b *PktRecordBucket) Get(key string) *PktRecord {
	b.Mu.Lock()
	defer b.Mu.Unlock()

	return b.PktRecoredMap[key]
}
